#!/usr/bin/python3
import os
import cv2
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.externals import joblib

class SK_process(object):

    def read_data(self):
        # 读取数据
        data = []
        path = '../static/learn_img/'
        folder_list = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
                       'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
                       '云', '京', '冀', '吉', '宁', '川', '新', '晋', '桂', '沪', '津', '浙', '渝', '湘', '琼', '甘', '皖', '粤', '苏', '蒙', '藏', '豫', '贵', '赣', '辽', '鄂', '闽', '陕', '青', '鲁', '黑']
        n = 0
        # 目标值
        # target = []
        # data = []
        target = []
        data = []
        for i in folder_list:
            # 遍历文件夹
            folder = '%s%s' % (path, i)
            file_list = os.listdir(folder)
            target_list = [0 for i in folder_list]
            target_list[n] = 1
            for j in range(15):
                # 遍历文件
                # img = cv2.imread('%s/%s' % (folder, file))
                img = cv2.imdecode(np.fromfile('%s/%s' % (folder, file_list[j]), dtype=np.uint8), -1)
                img = self.change_img(img)
                data.append(img)
                target.append(n)
            n += 1


        return np.array(data[:500][:]), np.array(target[:500])

    def change_img(self, img):
        # 转换灰度图
        # 统一大小 14*24
        img = cv2.resize(img, (14, 24))
        # img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        img = cv2.threshold(img, 125, 1, cv2.THRESH_BINARY)
        img_l = []
        for i in img[1]:
            img_l.extend(i)
        return img_l

    def process(self):
        data, target = self.read_data()
        print(data.shape)
        print(target.shape)
        x_train, x_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=22)
        # 实例化
        estimator = KNeighborsClassifier(n_neighbors=5)
        print(x_train.shape)
        print(y_train.shape)
        estimator.fit(x_train, y_train)
        # 模型评估
        # 直接计算准确率
        score = estimator.score(x_test, y_test)
        print("准确率为：\n", score)
        # 保存结果
        joblib.dump(estimator, "test.pkl")

    def run(self):
        self.process()


if __name__ == '__main__':
    sk = SK_process()
    sk.run()

